# 读取files目录下的所有文件，进行OCR识别，将识别结果输出到output目录的一个文本文件中：文件名.后缀名.txt 

from paddleocr import PaddleOCR
import os
from colorama import Fore

# 初始化 PaddleOCR 模型
ocr = PaddleOCR(use_angle_cls=True, lang="ch")  

# 指定输入和输出目录
input_dir = './files'   # 输入目录
output_dir = './output'  # 输出目录

# 确保输出目录存在，如果不存在则创建
os.makedirs(output_dir, exist_ok=True)

# 获取输出目录中已存在的文本文件列表
existing_files = set(os.listdir(output_dir))

# 统计文件的数量
total_files = 0
processed_files = 0
skipped_files = 0

# 遍历输入目录中的所有文件
for filename in os.listdir(input_dir):
    
    # 拼接输入文件的完整路径
    input_file = os.path.join(input_dir, filename)

    # 判断当前路径是否为文件而非目录
    if os.path.isfile(input_file):
        
        # 更新已检查的文件数量
        total_files += 1
        
        # 指定OCR支持文件格式后缀
        supported_formats = ['.png', '.jpg', '.jpeg', '.pdf']
        
        # 获取文件后缀
        ext = os.path.splitext(filename)[1]
        
        # 判断文件扩展名是否为支持 OCR 的格式
        if ext.lower() in supported_formats:
            # 构造输出文件名 文件名.后缀名.txt
            output_filename = filename + ".txt"
            
            # 设置输出文件路径
            output_file = os.path.join(output_dir, output_filename)
            
            # 检查是否已存在相应的文本文件，如果存在说明已经识别过，跳过当前文件
            if output_filename in existing_files:
                # 更新跳过的文件数量
                skipped_files += 1
                # 打印红色提示
                print(Fore.RED + f"==> 跳过文件 {filename}: 文件已经识别")
                continue
            
            # 打印绿色提示
            print(Fore.BLUE +f"==> 开始识别文件 {filename}")

            # 使用 PaddleOCR 模型进行文字识别
            result = ocr.ocr(input_file, cls=True)
            
            # 更新已识别的文件数量
            processed_files += 1

            # 将识别结果拼接成一个字符串
            text_result = ""
            for idx in range(len(result)):
                res = result[idx]
                for line in res:
                    text_result += line[1][0] + "\n"

            # 将识别结果写入到文本文件中
            with open(output_file, 'w', encoding='utf-8') as f:
                f.write(text_result)
            # 打印绿色提示
            print(Fore.GREEN +f"==> {filename} 识别完成，识别结果已写入 {output_file}")
        else:
            # 更新跳过的文件数量
            skipped_files += 1
            # 打印红色提示
            print(Fore.RED + f"==> 跳过文件 {filename}: 类型不支持")

# 输出程序执行情况
print(Fore.YELLOW + f"==> 文件总数： {total_files} ，识别数：{processed_files} ，跳过：{skipped_files} ")


